import nibabel as nib
import os
import numpy as np
from pathlib import Path

count = 10
resultPath = "../result"
# path = "../data"
# filePathList = os.listdir(path)
path = os.path.normpath(r'../data')
filePathList = Path(path).rglob('*.nii.gz')

labelFileSet = set()
imageFileSet = set()
for filePath in filePathList:
    filePath = str(filePath)
    if "image" in filePath:
        imageFileSet.add(filePath)
    elif "label" in filePath:
        labelFileSet.add(filePath)
    else:
        print(f"文件{filePath}名称不规范，因为他不包含image或者label！！！")
        exit()
countNum = 0
for labelFile in labelFileSet:
    for imageFile in imageFileSet:
        if labelFile.split("_")[1] in imageFile:
            break

    labelNii = nib.load(labelFile)
    imageNii = nib.load(imageFile)
    label = labelNii.get_fdata()
    image = imageNii.get_fdata()
    new_label = []
    new_image = []
    # 记录符合需求的帧的编号
    for i in range(label.shape[-1]):
        tempResult = label[:, :, i]
        if np.sum(label[:, :, i] > 0) > count:
            # if np.sum(label[:, :, i] > 0) > count:
            # new_label.append(label[:, :, i])
            # new_image.append(image[:, :, i])

            # （1,x,y）
            new_image = nib.Nifti1Image(np.array([image[:, :, i]]).transpose((-1, 1, 0)), np.eye(4))
            new_label = nib.Nifti1Image(np.array([label[:, :, i]]).transpose((-1, 1, 0)), np.eye(4))

            pathName = os.path.join(resultPath)
            if not os.path.exists(pathName):
                os.makedirs(pathName)
            name = labelFile.split('\\')[labelFile.split('\\').__len__() - 1].split('_')
            nib.save(new_label, os.path.join(pathName, 'label', f'{name[0]}_{"%04d" % countNum}_{name[2]}'))
            nib.save(new_image, os.path.join(pathName, 'image', f'{name[0]}_{"%04d" % countNum}_{name[2]}'))
            countNum += 1
            print(pathName,imageFile)
